import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision import datasets

def main():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]
    )

    ##读入训练集，40张一组进行迭代
    train_dir = "./data/training_data"
    train_dataset = datasets.ImageFolder(train_dir, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=40,
                                               shuffle=True, num_workers=2)


    ##读入验证集，40张一组进行迭代
    vertification_dir = "./data/verification_data"
    vertification_dataset = datasets.ImageFolder(vertification_dir, transform=transform)
    vertification_loader = torch.utils.data.DataLoader(vertification_dataset, batch_size=40,
                                               shuffle=True, num_workers=2)

    vertification_data_iter = iter(vertification_loader)
    vertification_image, vertification_label = vertification_data_iter.next()


    net = LeNet()
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.0001)

    for epoch in range(5):

        running_loss = 0.0
        for step, data in enumerate(train_loader, start=0):
            inputs, labels =data

            optimizer.zero_grad()

            outputs = net(inputs)

            loss = loss_function(outputs,labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if step % 200 == 199:
                with torch.no_grad():
                    outputs = net(vertification_image)
                    predict_y = torch.max(outputs, dim=1)[1]
                    accuracy = torch.eq(predict_y, vertification_label).sum().item() / vertification_label.size(0)

                    print('[%d, %5d] train loss: %.3f test accuracy: %.3f' %
                          (epoch + 1, step + 1, running_loss /200, accuracy))
                    running_loss = 0.0

    print('Finished Training')

    save_path = './LeNet.pth'
    torch.save(net.state_dict(), save_path)


if __name__ == '__main__':
        main()